//
//  Sample.hpp
//  RayTracer
//
//  Created by 俞云康 on 2/17/16.
//  Copyright © 2016 yushroom. All rights reserved.
//

#ifndef Sample_hpp
#define Sample_hpp

#include "Math.hpp"
#include <vector>
#include <algorithm>

// sample is a pair of uniform variable (u1, u2), both in range [0, 1]
typedef Point2f Sample;

struct CameraSample {
    Point2f pFilm;
    Point2f pLens;
    float time;
};

class Sampler {
public:
    virtual ~Sampler();
    Sampler (int samplersPerPixel);
    virtual void startPixel(const Point2i &p);
    virtual float get1D() = 0;
    virtual Point2f get2D() = 0;
    CameraSample getCameraSample(const Point2i &pRaster);
    void request1DArray(int n);
    void request2DArray(int n);
    virtual int roundCount(int n) const { return n; }
    const float *   get1DArray(int n);
    const Point2f * get2DArray(int n);
    virtual bool startNextSample();
    virtual unique_ptr<Sampler> clone(int seed) = 0;
    virtual bool setSampleNumber(int sampleNumber);
    
    const int m_samplersPerPixel;
    
protected:
    Point2i     m_currentPixel;
    int         m_currentPixelSampleIndex;
    vector<int> m_samples1DArraySizes;
    vector<int> m_samples2DArraySizes;
    vector<vector<float>>   m_sampleArray1D;
    vector<vector<Point2f>> m_sampleArray2D;
private:
    size_t      m_array1DOffest;
    size_t      m_array2DOffset;
};

//class GlobalSampler : public Sampler {
//public:
//    GlobalSampler(int samplesPerPixel) : Sampler(samplesPerPixel) {}
//    
//    virtual void startPixel(const Point2i &p) override;
//    virtual float get1D() override;
//    virtual Point2f get2D() override;
//    virtual bool startNextSample() override;
//    virtual bool setSampleNumber(int sampleNumber) override;
//    
//    virtual int getIndexForSample(int sampleNum) const = 0;
//    virtual float sampleDimension(int index, int dimension) const = 0;
//
//private:
//    int m_dimension;
//    int m_intervalSampleIndex;
//    static const int arrayStartDim = 5;
//    int m_arrayEndDim;
//};

class RandomSampler : public Sampler {
public:
    RandomSampler(int samplesPerPixel, int seed = 0);
    virtual void startPixel(const Point2i &) override;
    virtual float get1D() override;
    virtual Point2f get2D() override;
    virtual unique_ptr<Sampler> clone(int seed) override;
    
private:
    RNG rng;
};


//class Sampler {
//public:
//    virtual void setFilmResolution(const float w, const float h, bool blocked) {};
//    virtual void generate(const int offsetX, const int offsetY) {
//        m_sampleIndex = 0;
//        m_dimension1DArray = m_dimension2DArray = 0;
//    }
//    virtual void advance() {
//        m_sampleIndex++;
//        m_dimension1DArray = m_dimension1DArray = 0;
//    }
//    virtual void setSampleIndex(size_t sampleIndex) {
//        m_sampleIndex = sampleIndex;
//        m_dimension1DArray = m_dimension2DArray = 0;
//    }
//    
//    virtual float next1D() = 0;
//    virtual float2 next2D() = 0;
//    
//    float2 *next2DArray(size_t size) {
//        Assert(m_req2D[m_dimension2DArray] == size);
//        if (m_dimension2DArray < m_req2D.size()) {
//            Assert(m_req2D[m_dimension2DArray] == size);
//            return m_sampleArrays2D[m_dimension2DArray++] + m_sampleIndex * size;
//        } else {
//            error("");
//            return nullptr;
//        }
//    }
//    
//    float *next1DArray(size_t size) {
//        Assert(m_req1D[m_dimension2DArray] == size);
//        if (m_dimension1DArray < m_req1D.size()) {
//            Assert(m_req1D[m_dimension1DArray] == size);
//            return m_sampleArrays1D[m_dimension1DArray++] + m_sampleIndex * size;
//        } else {
//            error("");
//            return nullptr;
//        }
//    }
//    
//    virtual void request1DArray(size_t size) {
//        m_req1D.push_back(size);
//        m_sampleArrays1D.push_back(new float[m_sampleCount * size]);
//    }
//    virtual void request2DArray(size_t size) {
//        m_req2D.push_back(size);
//        m_sampleArrays2D.push_back(new float2[m_sampleCount * size]);
//    }
//    inline size_t getSampleCount() const { return m_sampleCount; }
//    inline size_t getSampleIndex() const { return m_sampleIndex; }
//
//protected:
//    Sampler() {}
//    virtual ~Sampler() {
//        for (size_t i = 0; i < m_sampleArrays1D.size(); i++) {
//            if (m_sampleArrays1D[i] != nullptr)
//                delete [] m_sampleArrays1D[i];
//        }
//        for (size_t i = 0; i < m_sampleArrays2D.size(); i++) {
//            if (m_sampleArrays2D[i] != nullptr)
//                delete [] m_sampleArrays2D[i];
//        }
//    }
//    
//protected:
//    size_t m_sampleCount = 0;
//    size_t m_sampleIndex = 0;
//    vector<size_t> m_req1D, m_req2D;
//    vector<float *>  m_sampleArrays1D;
//    vector<float2 *> m_sampleArrays2D;
//    size_t m_dimension1DArray, m_dimension2DArray;
//};
//
//class RandomSampler : public Sampler {
//public:
//    RandomSampler() : Sampler(), m_RNG(new RNG()) {
//        m_sampleCount = 4;
//    }
//    
//    virtual void generate(const int offsetX, const int offsetY) override {
//        for (size_t i = 0; i < m_req1D.size(); ++i) {
//            for (size_t j = 0; j < m_sampleCount * m_req1D[i]; ++j)
//                m_sampleArrays1D[i][j] = m_RNG->RandomFloat();
//        }
//        for (size_t i = 0; i < m_req2D.size(); ++i) {
//            for (size_t j = 0; j < m_sampleCount * m_req2D[i]; ++j)
//                m_sampleArrays2D[i][j] = float2{m_RNG->RandomFloat(), m_RNG->RandomFloat()};
//        }
//        m_sampleIndex = 0;
//        m_dimension1DArray = m_dimension2DArray = 0;
//    }
//    
//    virtual float next1D() override {
//        return m_RNG->RandomFloat();
//    }
//    
//    virtual float2 next2D() override {
//        float value1 = m_RNG->RandomFloat();
//        float value2 = m_RNG->RandomFloat();
//        return {value1, value2};
//    }
//    
//private:
//    unique_ptr<RNG> m_RNG;
//};

//class Sampler {
//public:
//    size_t num_samples;					// the number of sample points in a pattern
//	size_t num_sets;					// the number of sample sets(patterns) stored
//    std::vector<Sample> samples;		// sample points on a unit square
//    std::vector<int> shuffled_indices;	// shuffled samples array indices;
//	size_t count = 0;					// the current number of sample points used
//    size_t jump = 0;					// random index jump
//    
//public:
//	Sampler(size_t num_samples, size_t num_sets = 83)
//    : num_samples(num_samples), num_sets(num_sets), samples(num_samples*num_sets), shuffled_indices(num_samples*num_sets) {
//        setup_shuffled_indices();
//    }
//
//    virtual ~Sampler() {};
//    
//    // generate sample patterns in a unit square
//    virtual void generate_samples(RNG& rng) = 0;
//    
//    // get next sample on unit square;
//    virtual Sample sample_unit_sqrare(RNG& rng) {
//        //return (samples[count++ % (num_samples * num_sets)]);
//        if (count % num_samples == 0)   // start of a new pixel
//            jump = (rng.RandomUInt() % num_sets) * num_samples;
//        return (samples[jump + shuffled_indices[jump + count++%num_samples]]);
//    }
//    
//    void setup_shuffled_indices() {
//        //shuffled_indices.clear();
//        std::vector<int> indices(num_samples);
//        
//        for (int i = 0; i < num_samples; ++i)
//            indices[i] = i;
//        for (int p = 0; p < num_sets; ++p) {
//            std::random_shuffle(indices.begin(), indices.end());
//            for (int j = 0; j < num_samples; ++j)
//                //shuffled_indices.push_back(indices[j]);
//                shuffled_indices[j + p * num_samples] = indices[j];
//        }
//    }
//};
//
//
//class RandomSampler : public Sampler {
//public:
//    RandomSampler(int n_samples, int num_sets = 83) : Sampler(n_samples, num_sets) {}
//    
//    virtual void generate_samples(RNG& rng) override {
//        for (int i = 0; i < num_samples; ++i) {
//            samples[i].x = rng.RandomFloat();
//            samples[i].y = rng.RandomFloat();
//        }
//    }
//};
//
//
//class JitterSampler : public Sampler {
//private:
//    size_t sqrt_n_samples;
//	float stratum_size;
//public:
//	JitterSampler(const size_t sqrt_n_samples, const size_t num_sets = 83)
//		: Sampler(sqrt_n_samples*sqrt_n_samples, num_sets), sqrt_n_samples(sqrt_n_samples), stratum_size(1.0f /(float)sqrt_n_samples) {
//        
//    }
//    
//    virtual void generate_samples(RNG& rng) override {
//        for (int p = 0; p < num_sets; p++) {
//            for (int y = 0; y < sqrt_n_samples; ++y) {
//                for (int x = 0; x < sqrt_n_samples; ++x) {
//                    auto& sa = samples[x+y*sqrt_n_samples+p*num_samples];
//                    sa.x = (x + rng.RandomFloat()) * stratum_size;
//                    sa.y = (y + rng.RandomFloat()) * stratum_size;
//                }
//            }
//        }
//    }
//};


// Sampling
inline float3 UniformSampleHemisphere(float u1, float u2) {
    float z = u1;
    float r = sqrtf(fmax(0.f, 1.f - z*z));
    float phi = 2 * M_PI * u2;
    float x = r * cosf(phi);
    float y = r * sinf(phi);
    return float3(x, y, z);
}


inline constexpr float UniformHemispherePdf()
{
    return 1.0f / (2 * M_PI);
}


inline float3 UniformSampleSphere(float u1, float u2) {
    float z = 1.f - 2.f * u1;
    float r = sqrtf(fmax(0.f, 1.f - z*z));
    float phi = 2.f * M_PI * u2;
    float x = r * cosf(phi);
    float y = r * sinf(phi);
    return float3(x, y, z);
}


inline constexpr float UniformSpherePdf()
{
    return 1.0f / (4.f * M_PI);
}


inline void UniformSampleTriangle(float u1, float u2, float *u, float *v) {
	float su1 = sqrtf(u1);
	*u = 1.f - su1;
	*v = u2 * su1;
}


inline float VanDerCorput(uint32_t n, uint32_t scramble) {
    // Reverse bits of _n_
    n = (n << 16) | (n >> 16);
    n = ((n & 0x00ff00ff) << 8) | ((n & 0xff00ff00) >> 8);
    n = ((n & 0x0f0f0f0f) << 4) | ((n & 0xf0f0f0f0) >> 4);
    n = ((n & 0x33333333) << 2) | ((n & 0xcccccccc) >> 2);
    n = ((n & 0x55555555) << 1) | ((n & 0xaaaaaaaa) >> 1);
    n ^= scramble;
    return fmin(((n>>8) & 0xffffff) / float(1 << 24), OneMinusEpsilon);
}


inline float Sobol2(uint32_t n, uint32_t scramble) {
    for (uint32_t v = 1 << 31; n != 0; n >>= 1, v ^= v >> 1)
        if (n & 0x1) scramble ^= v;
    return fmin(((scramble>>8) & 0xffffff) / float(1 << 24), OneMinusEpsilon);
}


inline void Sample02(uint32_t n, const uint32_t scramble[2], float sample[2])
{
    sample[0] = VanDerCorput(n, scramble[0]);
    sample[1] = Sobol2(n, scramble[1]);
}


// see PBRT 667
void ConcentricSampleDisk(float u1, float u2, float* dx, float* dy);


inline void ConcentricSampleDisk(const Sample& sample, float* dx, float* dy) {
    return ConcentricSampleDisk(sample.x, sample.y, dx, dy);
}


inline float UniformConePdf(float cosThetaMax) {
    return 1.f / (2.f * M_PI * (1.f - cosThetaMax));
}


inline Direction UniformSampleCone(float u1, float u2, float costhetamax) {
    float costheta = (1.f - u1) + u1 * costhetamax;
    float sintheta = sqrtf(1.f - costheta*costheta);
    float phi = u2 * 2.f * M_PI;
    return Direction(cosf(phi) * sintheta, sinf(phi) * sintheta, costheta);
}


inline Direction UniformSampleCone(float u1, float u2, float costhetamax,
                         const Direction &x, const Direction &y, const Direction &z) {
    float costheta = Lerp(u1, costhetamax, 1.f);
    float sintheta = sqrtf(1.f - costheta*costheta);
    float phi = u2 * 2.f * M_PI;
    return cosf(phi) * sintheta * x + sinf(phi) * sintheta * y +
    costheta * z;
}


inline Direction CosineSampleHemisphere(const Sample& sample) {
    Direction ret;
    ConcentricSampleDisk(sample, &ret.x, &ret.y);
    ret.z = sqrtf(std::max(0.f, 1.f - ret.x*ret.x - ret.y*ret.y));
    return ret;
}


inline bool SameHemisphere(const Direction& w, const Direction& wp) {
    return w.z * wp.z > 0.f;
}


// cos\theta = (n\dot\omega) = ((0, 0, 1)\dot\omega) = \omega_z
inline float CosTheta(const Direction& w) {
    return w.z;
}


inline float AbsCosTheta(const Direction& w) {
    return fabsf(w.z);
}


// sin\theta^2 = 1 - cos\theta^2
inline float SinTheta2(const Direction& w) {
	//float ret = 1.f - CosTheta(w)*CosTheta(w);
	//assert(ret >= 0);
	//return ret;
    return std::max(0.f, 1.f - CosTheta(w)*CosTheta(w));
}


inline float SinTheta(const Direction& w) {
    return sqrtf(SinTheta2(w));
}


inline float CosPhi(const Direction &w) {
    float sintheta = SinTheta(w);
    if (sintheta == 0.f) return 1.f;
    return clamp(w.x / sintheta, -1.f, 1.f);
}


inline float SinPhi(const Direction &w) {
    float sintheta = SinTheta(w);
    if (sintheta == 0.f) return 0.f;
    return clamp(w.y / sintheta, -1.f, 1.f);
}


inline Sample random_sample(RNG& rng) {
    return {rng.UniformFloat(), rng.UniformFloat()};
}

inline Direction SphericalDirection(float sintheta,
                                 float costheta, float phi) {
    return Direction(sintheta * cosf(phi),
                  sintheta * sinf(phi),
                  costheta);
}

inline Direction SphericalDirection(float sintheta, float costheta,
                                 float phi, const Direction &x,
                                 const Direction &y, const Direction &z) {
    return sintheta * cosf(phi) * x +
    sintheta * sinf(phi) * y + costheta * z;
}

inline float SphericalTheta(const Direction &v) {
    return acosf(clamp(v.z, -1.f, 1.f));
}

inline float SphericalPhi(const Direction &v) {
    float p = atan2f(v.y, v.x);
    return (p < 0.f) ? p + 2.f*M_PI : p;
}

#endif /* Sample_hpp */
